package org.rcy.framework.data.transaction;

import com.mysql.cj.jdbc.JdbcConnection;
import com.zaxxer.hikari.pool.HikariProxyConnection;
import com.zaxxer.hikari.pool.ProxyConnection;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.jdbc.datasource.DataSourceTransactionManager;
import org.springframework.jdbc.datasource.JdbcTransactionObjectSupport;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.support.DefaultTransactionStatus;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.sql.Connection;

/**
 * @author 吴康桥
 * @Description
 * @date 2022/9/28 9:43
 */


public class RcyDataSourceTransactionManager extends DataSourceTransactionManager {
	private static final long serialVersionUID = 1L;

	private static final Logger log = LoggerFactory.getLogger(RcyDataSourceTransactionManager.class);

	private static final ThreadLocal<TransactionInfo> TRANSACTION_THREAD = new ThreadLocal<>();

	private static Field cachedDelegateField;

	private RcyDataSourceContext dataSourceContext;

	static {
		try {
			cachedDelegateField = ProxyConnection.class.getDeclaredField("delegate");
			cachedDelegateField.setAccessible(true);
		} catch (NoSuchFieldException | SecurityException e) {
			log.error("===> Get cached delegate field error:{} ", e.getMessage());
		}
	}

	@SuppressWarnings("unused")
	private RcyDataSourceTransactionManager() {
		super();
	}

	public RcyDataSourceTransactionManager(DataSource dataSource, RcyDataSourceContext dataSourceContext) {
		super(dataSource);
		this.dataSourceContext = dataSourceContext;
		this.dataSourceContext.startSchedule();
	}

	@Override
	protected void doBegin(Object transaction, TransactionDefinition definition) {
		long startTime = System.currentTimeMillis();
		super.doBegin(transaction, definition);
		try {
			String connectionId = getConnectionId(transaction);
			String traceId = MDC.get("traceId");
			if(connectionId == null) {
				connectionId = traceId;
			}
			TransactionInfo transactionInfo = dataSourceContext.addTransactionInfo(definition.getName(), startTime, traceId, connectionId);
			if(transactionInfo == null) {
				return;
			}
			TRANSACTION_THREAD.set(transactionInfo);
			log.debug("transaction doBegin connectionId {} currentThread {}",connectionId,Thread.currentThread().getId());
		} catch (Exception e) {
			log.error("transaction doBegin error {}", e.getMessage(), e);
		}
	}

	public String getConnectionId(Object transaction){
		try {
			if(transaction instanceof JdbcTransactionObjectSupport) {
				JdbcTransactionObjectSupport jdbcTransactionObjectSupport = (JdbcTransactionObjectSupport) transaction;
				Connection connection = jdbcTransactionObjectSupport.getConnectionHolder().getConnection();
				if(connection instanceof HikariProxyConnection) {
					HikariProxyConnection proxyConnection = (HikariProxyConnection) connection;
					JdbcConnection jdbcConnection = (JdbcConnection) cachedDelegateField.get(proxyConnection);
					return String.valueOf(jdbcConnection.getId());
				}
			}
		} catch (Exception e) {
			log.warn("qys transaction getConnectionId error {}", e.getMessage());
		}
		return null;
	}

	@Override
	protected void doCommit(DefaultTransactionStatus status) {
		super.doCommit(status);
		calcTransactionTime("doCommit");
	}

	@Override
	protected void doRollback(DefaultTransactionStatus status) {
		super.doRollback(status);
		calcTransactionTime("doRollback");
	}

	private void calcTransactionTime(String type) {
		try {
			TransactionInfo transactionInfo = TRANSACTION_THREAD.get();
			if (transactionInfo == null) {
				return;
			}
			long transactionCost = transactionInfo.cost(System.currentTimeMillis());
			String connectionId = transactionInfo.getConnectionId();
			if (transactionCost > dataSourceContext.getMaxTransactionTime()) {
				log.warn("transaction method {} connectionId {} {} cost {}ms",transactionInfo.getMethod(), connectionId ,type, transactionCost);
			}
			dataSourceContext.removeTransactionInfo(connectionId);
		} catch (Exception e) {
			log.error("transaction calcTime error {}", e.getMessage(), e);
		}
	}

	@Override
	protected void doCleanupAfterCompletion(Object transaction) {
		super.doCleanupAfterCompletion(transaction);
		log.debug("transaction threadLocal remove currentThread {}",Thread.currentThread().getId());
		TRANSACTION_THREAD.remove();
	}

}
